import ctypes
import sys
from ctypes import (c_void_p, c_char_p, c_size_t, c_int, c_uint,
                    POINTER, Structure, create_string_buffer)

# 加载 HIP 库
try:
    libgalaxyhip = ctypes.CDLL("libgalaxyhip.so")
    libhiprtc = ctypes.CDLL("libhiprtc.so")
except Exception as e:
    print(f"Failed to load HIP libraries: {e}")
    sys.exit(1)


# 定义 HIP 结构体和常量
class hipDeviceProp_t(Structure):
    _fields_ = [
        ("name", ctypes.c_char * 256),
        ("totalGlobalMem", ctypes.c_size_t),
        ("sharedMemPerBlock", ctypes.c_size_t),
    ]


HIPRTC_SUCCESS = 0
HIP_MEMCPY_HOST_TO_DEVICE = 1
HIP_MEMCPY_DEVICE_TO_HOST = 2

# 定义 HIP Runtime Compilation API
libhiprtc.hiprtcCreateProgram.restype = c_int
libhiprtc.hiprtcCreateProgram.argtypes = [POINTER(c_void_p), c_char_p, c_char_p, c_int, POINTER(c_char_p),
                                          POINTER(c_char_p)]

libhiprtc.hiprtcCompileProgram.restype = c_int
libhiprtc.hiprtcCompileProgram.argtypes = [c_void_p, c_int, POINTER(c_char_p)]

libhiprtc.hiprtcGetCodeSize.restype = c_int
libhiprtc.hiprtcGetCodeSize.argtypes = [c_void_p, POINTER(c_size_t)]

libhiprtc.hiprtcGetCode.restype = c_int
libhiprtc.hiprtcGetCode.argtypes = [c_void_p, c_char_p]

libhiprtc.hiprtcDestroyProgram.restype = c_int
libhiprtc.hiprtcDestroyProgram.argtypes = [POINTER(c_void_p)]

# 添加缺失的 hiprtc 函数
libhiprtc.hiprtcGetProgramLogSize.restype = c_int
libhiprtc.hiprtcGetProgramLogSize.argtypes = [c_void_p, POINTER(c_size_t)]

libhiprtc.hiprtcGetProgramLog.restype = c_int
libhiprtc.hiprtcGetProgramLog.argtypes = [c_void_p, c_char_p]

# 定义 HIP API 函数
hipGetDeviceProperties = libgalaxyhip.hipGetDeviceProperties
hipGetDeviceProperties.argtypes = [POINTER(hipDeviceProp_t), c_int]
hipGetDeviceProperties.restype = c_int

hipModuleLoadData = libgalaxyhip.hipModuleLoadData
hipModuleLoadData.argtypes = [POINTER(c_void_p), c_char_p]
hipModuleLoadData.restype = c_int

hipModuleGetFunction = libgalaxyhip.hipModuleGetFunction
hipModuleGetFunction.argtypes = [POINTER(c_void_p), c_void_p, c_char_p]
hipModuleGetFunction.restype = c_int

hipMalloc = libgalaxyhip.hipMalloc
hipMalloc.argtypes = [POINTER(c_void_p), c_size_t]
hipMalloc.restype = c_int

hipMemcpy = libgalaxyhip.hipMemcpy
hipMemcpy.argtypes = [c_void_p, c_void_p, c_size_t, c_int]
hipMemcpy.restype = c_int

hipFree = libgalaxyhip.hipFree
hipFree.argtypes = [c_void_p]
hipFree.restype = c_int

hipModuleUnload = libgalaxyhip.hipModuleUnload
hipModuleUnload.argtypes = [c_void_p]
hipModuleUnload.restype = c_int

# 添加 hipDeviceSynchronize
hipDeviceSynchronize = libgalaxyhip.hipDeviceSynchronize
hipDeviceSynchronize.argtypes = []
hipDeviceSynchronize.restype = c_int

# 添加 hipGetLastError
hipGetLastError = libgalaxyhip.hipGetLastError
hipGetLastError.argtypes = []
hipGetLastError.restype = c_int

# 定义 hipModuleLaunchKernel
hipModuleLaunchKernel = libgalaxyhip.hipModuleLaunchKernel
hipModuleLaunchKernel.argtypes = [
    c_void_p,  # function
    c_uint, c_uint, c_uint,  # gridDimX, gridDimY, gridDimZ
    c_uint, c_uint, c_uint,  # blockDimX, blockDimY, blockDimZ
    c_uint,  # sharedMemBytes
    c_void_p,  # stream
    c_void_p,  # kernelParams
    c_void_p,  # extra
]
hipModuleLaunchKernel.restype = c_int


class HIPContext:
    def __init__(self):
        self.prop = hipDeviceProp_t()
        self.module = c_void_p()
        self.kernel_func = c_void_p()

        # 初始化设备属性
        if hipGetDeviceProperties(ctypes.byref(self.prop), 0) != 0:
            raise RuntimeError("Failed to get device properties")
        print(f"Device: {self.prop.name.decode()}")
        print(f"Total Global Memory: {self.prop.totalGlobalMem} bytes")

        # 动态编译核函数
        self._compile_kernel()

    def _compile_kernel(self):
        kernel_code = """
extern "C" __global__ void matrixTranspose(float* out, float* in, int width) {
    int x = blockIdx.x * blockDim.x + threadIdx.x;
    int y = blockIdx.y * blockDim.y + threadIdx.y;

    if (x < width && y < width) {
        out[y * width + x] = in[x * width + y];
    }
}
"""
        prog = c_void_p()
        options = [b"-O3"]

        print("Creating HIPRTC program...")
        result = libhiprtc.hiprtcCreateProgram(ctypes.byref(prog),
                                               c_char_p(kernel_code.encode()),
                                               b"matrixTranspose.cu",
                                               0,
                                               None,
                                               None)
        if result != HIPRTC_SUCCESS:
            self._handle_hiprtc_error(prog, f"Failed to create HIPRTC program: {result}")

        print("Compiling kernel...")
        opts = (c_char_p * len(options))(*options)
        result = libhiprtc.hiprtcCompileProgram(prog, len(options), opts)
        if result != HIPRTC_SUCCESS:
            self._handle_hiprtc_error(prog, f"Kernel compilation failed: {result}")

        print("Getting compiled code...")
        code_size = c_size_t()
        if libhiprtc.hiprtcGetCodeSize(prog, ctypes.byref(code_size)) != HIPRTC_SUCCESS:
            self._handle_hiprtc_error(prog, "Failed to get code size")

        code = create_string_buffer(code_size.value)
        if libhiprtc.hiprtcGetCode(prog, code) != HIPRTC_SUCCESS:
            self._handle_hiprtc_error(prog, "Failed to get compiled code")

        print("Loading module...")
        if hipModuleLoadData(ctypes.byref(self.module), code) != 0:
            raise RuntimeError("Failed to load module")

        print("Getting kernel function...")
        if hipModuleGetFunction(ctypes.byref(self.kernel_func),
                                self.module,
                                b"matrixTranspose") != 0:
            raise RuntimeError("Failed to get kernel function")

        libhiprtc.hiprtcDestroyProgram(ctypes.byref(prog))
        print("Kernel compilation successful")

    def _handle_hiprtc_error(self, prog, msg):
        log_size = c_size_t()
        if libhiprtc.hiprtcGetProgramLogSize(prog, ctypes.byref(log_size)) == HIPRTC_SUCCESS:
            log = create_string_buffer(log_size.value)
            if libhiprtc.hiprtcGetProgramLog(prog, log) == HIPRTC_SUCCESS:
                log_text = log.value.decode()
                print("Compilation log:")
                print(log_text)
                msg += f"\nCompilation log:\n{log_text}"
        raise RuntimeError(msg)

    def __del__(self):
        if self.module:
            hipModuleUnload(self.module)


class HIPMemory:
    def __init__(self, size):
        self.ptr = c_void_p()
        if hipMalloc(ctypes.byref(self.ptr), size) != 0:
            raise RuntimeError("Failed to allocate HIP memory")
        print(f"Allocated {size} bytes of device memory")

    def __del__(self):
        if self.ptr:
            hipFree(self.ptr)


def main():
    try:
        ctx = HIPContext()

        WIDTH = 256
        NUM = WIDTH * WIDTH
        SIZE = NUM * ctypes.sizeof(ctypes.c_float)

        print(f"Matrix size: {WIDTH}x{WIDTH}")
        print(f"Total elements: {NUM}")
        print(f"Memory size: {SIZE} bytes")

        # 初始化数据
        host_input = (ctypes.c_float * NUM)()
        host_output = (ctypes.c_float * NUM)()
        for i in range(NUM):
            host_input[i] = float(i)

        print("Initialized host data")

        # 分配设备内存
        d_input = HIPMemory(SIZE)
        d_output = HIPMemory(SIZE)

        # 数据传输到设备
        print("Copying data to device...")
        if hipMemcpy(d_input.ptr, host_input, SIZE, HIP_MEMCPY_HOST_TO_DEVICE) != 0:
            raise RuntimeError("Failed to copy data to device")

        if hipDeviceSynchronize() != 0:
            raise RuntimeError("Failed to synchronize after memory copy to device")

        # 启动核函数
        block_dim = (16, 16, 1)
        grid_dim = (
            (WIDTH + block_dim[0] - 1) // block_dim[0],
            (WIDTH + block_dim[1] - 1) // block_dim[1],
            1
        )

        print(f"Launching kernel with grid: {grid_dim}, block: {block_dim}")

        # 准备核函数参数
        args = [ctypes.cast(d_output.ptr, c_void_p),
                ctypes.cast(d_input.ptr, c_void_p),
                c_int(WIDTH)]

        # 创建参数指针数组
        arg_ptrs = (c_void_p * len(args))()
        for i, arg in enumerate(args):
            if hasattr(arg, '_type_') and hasattr(arg, 'value'):
                arg_ptrs[i] = ctypes.cast(ctypes.byref(arg), c_void_p)
            else:
                arg_ptrs[i] = ctypes.cast(arg, c_void_p)

        print(f"Kernel arguments prepared: {len(args)} parameters")

        # 启动核函数
        print("Launching kernel...")
        result = hipModuleLaunchKernel(
            ctx.kernel_func,
            c_uint(grid_dim[0]), c_uint(grid_dim[1]), c_uint(grid_dim[2]),
            c_uint(block_dim[0]), c_uint(block_dim[1]), c_uint(block_dim[2]),
            c_uint(0),  # shared memory
            None,  # stream
            ctypes.cast(arg_ptrs, c_void_p),  # kernel parameters
            None  # extra
        )

        if result != 0:
            error_code = hipGetLastError()
            raise RuntimeError(f"Failed to launch kernel: {result}, last error: {error_code}")

        print("Kernel launched successfully")

        # 等待核函数完成
        if hipDeviceSynchronize() != 0:
            raise RuntimeError("Kernel execution failed or synchronization error")

        print("Kernel execution completed successfully")

        # 回传数据到主机
        print("Copying data back to host...")
        if hipMemcpy(host_output, d_output.ptr, SIZE, HIP_MEMCPY_DEVICE_TO_HOST) != 0:
            raise RuntimeError("Failed to copy data from device")

        if hipDeviceSynchronize() != 0:
            raise RuntimeError("Failed to synchronize after memory copy from device")

        # 验证结果
        print("Validating results...")
        errors = 0
        for j in range(min(10, WIDTH)):
            for i in range(min(10, WIDTH)):
                expected = host_input[j * WIDTH + i]
                actual = host_output[i * WIDTH + j]
                if abs(actual - expected) > 1e-6:
                    errors += 1
                    if errors <= 5:
                        print(f"Error at ({i},{j}): expected {expected}, got {actual}")

        if errors == 0:
            print("Validation: PASSED")
        else:
            print(f"Validation: FAILED ({errors} errors found in sampled data)")

    except Exception as e:
        print(f"Error: {str(e)}")
        import traceback
        traceback.print_exc()
        sys.exit(1)


if __name__ == "__main__":
    main()